#include <stdio.h>
#include <stdlib.h>
#include <time.h>
#include <arm_sve.h>
#include <string.h>
#define SIZE (1 << 16)  // 测试数据大小：16MB
#define ITERATIONS 1000000   // 重复次数

void scalar_memcpy(char *dst, const char *src, size_t n) {
    for (size_t i = 0; i < n; i+=256) {
        asm volatile(
            "ldp q0, q1, [%[src]]\n"
            "ldp q2, q3, [%[src], 32]\n"
            "ldp q4, q5, [%[src], 64]\n"
            "ldp q6, q7, [%[src], 96]\n"
            "ldp q8, q9, [%[src], 128]\n"
            "ldp q10, q11, [%[src], 160]\n"
            "ldp q12, q13, [%[src], 192]\n"
            "ldp q14, q15, [%[src], 224]\n"
	    "stp q0, q1, [%[dst]]\n"
	    "stp q2, q3, [%[dst], 32]\n"
	    "stp q4, q5, [%[dst], 64]\n"
	    "stp q6, q7, [%[dst], 96]\n"
	    "stp q8, q9, [%[dst], 128]\n"
	    "stp q10, q11, [%[dst], 160]\n"
	    "stp q12, q13, [%[dst], 192]\n"
	    "stp q14, q15, [%[dst], 224]\n"
            : 
            : [src] "r" (&src[i]), [dst] "r" (&dst[i])
            : "q0", "q1", "q2", "q3", "q4", "q5", "q6", "q7", "q8", "q9", "q10", "q11", "q12", "q13", "q14", "q15", "memory"
        );
    }
}

void sve_memcpy(char *dst, const char *src, size_t n) {
    for (size_t i = 0; i < n; i+=256) {
        asm volatile(
	    "ptrue p0.b\n"
            "ld1b {z0.b}, p0/z, [%[src]]\n"
            "ld1b {z1.b}, p0/z, [%[src], 1, mul vl]\n"
            "ld1b {z2.b}, p0/z, [%[src], 2, mul vl]\n"
            "ld1b {z3.b}, p0/z, [%[src], 3, mul vl]\n"
            "ld1b {z4.b}, p0/z, [%[src], 4, mul vl]\n"
            "ld1b {z5.b}, p0/z, [%[src], 5, mul vl]\n"
            "ld1b {z6.b}, p0/z, [%[src], 6, mul vl]\n"
            "ld1b {z7.b}, p0/z, [%[src], 7, mul vl]\n"
            "st1b {z0.b}, p0, [%[dst]]\n"
            "st1b {z1.b}, p0, [%[dst], 1, mul vl]\n"
            "st1b {z2.b}, p0, [%[dst], 2, mul vl]\n"
            "st1b {z3.b}, p0, [%[dst], 3, mul vl]\n"
            "st1b {z4.b}, p0, [%[dst], 4, mul vl]\n"
            "st1b {z5.b}, p0, [%[dst], 5, mul vl]\n"
            "st1b {z6.b}, p0, [%[dst], 6, mul vl]\n"
            "st1b {z7.b}, p0, [%[dst], 7, mul vl]\n"
            : 
            : [src] "r" (&src[i]), [dst] "r" (&dst[i])
            : "z0", "z1", "z2", "z3", "z4", "z5", "z6", "z7", "memory"
        );
    }
}

// 计时函数
double measure_time(void (*func)(char *, const char *, size_t),
                   char *dst, char *src, size_t n) {
    struct timespec start, end;
    clock_gettime(CLOCK_MONOTONIC, &start);
    for (int i = 0; i < ITERATIONS; i++) {
        func(dst, src, n);
    }
    clock_gettime(CLOCK_MONOTONIC, &end);
    double time = (end.tv_sec - start.tv_sec) * 1e9 + (end.tv_nsec - start.tv_nsec);
    return time / ITERATIONS;  // 平均时间（纳秒）
}

int main() {
    char *src = malloc(SIZE);
    char *dst = malloc(SIZE);

    // 初始化数据
    for (size_t i = 0; i < SIZE; i++) {
        src[i] = (char)(i % 256);
    }

    // 测试标量拷贝
    double scalar_time = measure_time(scalar_memcpy, dst, src, SIZE);
    double scalar_bandwidth = (SIZE / scalar_time) * 1e9 / (1024 * 1024 * 1024);  // GB/s

    // 测试 SVE 拷贝
    double sve_time = measure_time(sve_memcpy, dst, src, SIZE);
    double sve_bandwidth = (SIZE / sve_time) * 1e9 / (1024 * 1024 * 1024);  // GB/s

    printf("Scalar memcpy: %.2f ns | Bandwidth: %.2f GB/s\n", scalar_time, scalar_bandwidth);
    printf("SVE memcpy:    %.2f ns | Bandwidth: %.2f GB/s\n", sve_time, sve_bandwidth);
    printf("Speedup: %.2fx\n", scalar_time / sve_time);

    free(src);
    free(dst);
    return 0;
}
